from flask import Flask, render_template, request, redirect, url_for, flash, send_file
import pandas as pd
import os
import io
import uuid
from datetime import datetime

app = Flask(__name__)
app.config['SECRET_KEY'] = 'your-secret-key-here'  # 用于flash消息
app.config['UPLOAD_FOLDER'] = 'uploads'
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 限制上传文件大小为16MB

# 确保上传目录存在
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)

# 允许的文件扩展名
ALLOWED_EXTENSIONS = {'xlsx', 'xls'}

# 用于临时存储处理后的文件数据（生产环境可考虑使用数据库或缓存）
processed_files = {}

def allowed_file(filename):
    """检查文件是否为允许的类型"""
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

def process_file(file_stream, time_column='B'):
    """处理上传的文件，添加时段标记并生成统计数据"""
    try:
        # 解决文件流不可seek问题
        file_content = file_stream.read()
        bytes_io = io.BytesIO(file_content)
        bytes_io.seek(0)
        
        # 读取Excel文件
        df = pd.read_excel(bytes_io)
        
        # 将列标转为列名（获取时间列的实际列名）
        time_col_index = ord(time_column.upper()) - ord('A')
        if time_col_index < 0 or time_col_index >= len(df.columns):
            raise ValueError(f"指定的时间列 {time_column} 不存在于文件中")
            
        time_col_name = df.columns[time_col_index]
        
        # 转换时间格式
        df[time_col_name] = pd.to_datetime(df[time_col_name], errors='coerce')
        
        # 添加时段标记列
        df['所属2小时区间'] = df[time_col_name].apply(lambda x: 
            f"{(x.hour // 2) * 2:02d}:00-{(x.hour // 2) * 2 + 2:02d}:00" 
            if pd.notna(x) else "无效时间"
        )
        
        # 生成统计数据
        stats = df.groupby('所属2小时区间').size().reset_index(name='数据量')
        
        # 确保统计结果按时间顺序排列
        time_order = [f"{h:02d}:00-{h+2:02d}:00" for h in range(0, 24, 2)] + ["无效时间"]
        stats['所属2小时区间'] = pd.Categorical(stats['所属2小时区间'], categories=time_order, ordered=True)
        stats = stats.sort_values('所属2小时区间').reset_index(drop=True)
        
        return df, stats, None
        
    except Exception as e:
        return None, None, str(e)

@app.route('/', methods=['GET', 'POST'])
def upload_file():
    if request.method == 'POST':
        # 检查是否有文件部分
        if 'file' not in request.files:
            flash('没有文件部分')
            return redirect(request.url)
            
        file = request.files['file']
        
        # 检查用户是否选择了文件
        if file.filename == '':
            flash('没有选择文件')
            return redirect(request.url)
            
        # 检查文件是否有效
        if file and allowed_file(file.filename):
            # 获取用户指定的时间列
            time_column = request.form.get('time_column', 'B').strip() or 'B'
            
            # 处理文件
            df, stats, error = process_file(file.stream, time_column)
            
            if error:
                flash(f'处理文件时出错: {error}')
                return redirect(request.url)
            
            # 准备下载数据
            output = io.BytesIO()
            with pd.ExcelWriter(output, engine='openpyxl') as writer:
                df.to_excel(writer, sheet_name='标记后的数据', index=False)
                stats.to_excel(writer, sheet_name='统计结果', index=False)
            output.seek(0)
            
            # 生成唯一ID存储处理后的文件
            file_id = str(uuid.uuid4())
            processed_files[file_id] = output.getvalue()
            
            # 将结果转换为HTML表格
            df_html = df.to_html(classes='table table-striped', index=False)
            stats_html = stats.to_html(classes='table table-success', index=False)
            
            # 生成下载文件名
            timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
            original_name = os.path.splitext(file.filename)[0]
            download_filename = f"{original_name}_processed_{timestamp}.xlsx"
            
            return render_template('result.html', 
                                 df_html=df_html, 
                                 stats_html=stats_html,
                                 download_filename=download_filename,
                                 file_id=file_id)
    
    # GET请求，显示上传表单
    return render_template('upload.html')

@app.route('/download/<file_id>/<filename>')
def download_file(file_id, filename):
    # 检查文件是否存在
    if file_id not in processed_files:
        flash('文件不存在或已过期')
        return redirect(url_for('upload_file'))
    
    # 获取文件数据
    file_data = processed_files[file_id]
    output = io.BytesIO(file_data)
    output.seek(0)
    
    # 可选：删除已下载的文件数据释放内存
    # del processed_files[file_id]
    
    return send_file(
        output,
        mimetype='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
        as_attachment=True,
        download_name=filename
    )

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=False)